﻿#include "HeatConductionSystem.h"
#include "Mesh.h"
#include <vtkm/exec/CellDerivative.h>
#include <vtkm/exec/CellInterpolate.h>
#include <vtkm/worklet/WorkletMapField.h>
#include <vtkm/worklet/WorkletMapTopology.h>
#include <vtkm/cont/ArrayHandleGroupVec.h>


struct InitTemperature : vtkm::worklet::WorkletMapField
{
  using ControlSignature = void(FieldIn pts, FieldOut temperature);
  using ExecutionSignature = void(_1, _2);

  template<typename CoordType>
  VTKM_EXEC void operator()(CoordType& point, Real& value) const
  {
    value = 0.1;
  }
};

struct ComputeRhs : vtkm::worklet::WorkletVisitCellsWithPoints
{
  // clang-format off
  using ControlSignature =  void(CellSetIn, FieldInPoint, FieldInPoint temperature, FieldOut rhs, FieldOut c2n);
  // clang-format on
  using ExecutionSignature = void(_2, _3, _4, PointIndices, _5);

  template<typename PType, typename TempType, typename RhsType, typename C2NType>
  VTKM_EXEC void operator()(const PType& pts,
                            const TempType& temperature,
                            RhsType& rhs,
                            const C2NType& c2n,
                            Vec<Id, 3>& conn) const
  {
    vtkm::Vec<Real, 3> pcoord{ 0.5, 0.5, 0 };
    vtkm::Vec<Real, 3> result;
    vtkm::exec::CellDerivative(temperature, pts, pcoord, vtkm::CellShapeTagTriangle{}, result);
  }

  template<typename WorldCoordType, typename ParametricCoordType>
  VTKM_EXEC void GradPhi(vtkm::UInt8 i,
                         const WorldCoordType& wcoord,
                         const vtkm::Vec<ParametricCoordType, 3>& pcoord) const
  {
  }
};

RegisterObject(HeatConductionSystem);

HeatConductionSystem::HeatConductionSystem(const Configuration& cfg)
  : System(cfg)
{
}

void HeatConductionSystem::Init()
{
  _data_set.SetCellSet(_mesh.GetCellSet());
  _data_set.AddCoordinateSystem(_mesh.GetCoordinate());
  this->AddPointVariable<Real>("T");

  auto temperature = GetVariableArray<Real>("T");
  invoker(InitTemperature{}, _mesh.GetCoordinate(), temperature);
  console::Info("HeatConductionSystem::Init");
}

void HeatConductionSystem::InitVariable() {}

void HeatConductionSystem::Evolve()
{
  console::Info("HeatConductionSystem::Evolve");

  auto temperature = GetVariableArray<Real>("T");
  invoker(ComputeRhs{},
          _mesh.GetCellSet(),
          _mesh.GetCoordinate(),
          temperature,
          vtkm::cont::make_ArrayHandleGroupVec<3>(_c2n_rhs),
          vtkm::cont::make_ArrayHandleGroupVec<3>(_c2n));
}
